########################
# astar_fibheap.py
#
# An A* search-problem class utilizing a Fibonacci heap
# for the decrease-key operations.

# Subclass AStar for your specific problem domain!
#
# Written: Mark Wilson (mw54) 9/3/2010
# Based on code by Ikhyun Park
# Last updated: 9/6/2010
########################

"""This is the astar_fibheap module.  It defines classes for an A* search
problem and nodes in the search tree.  This module uses
the fibheap.FibonacciHeap type for the search's priority-queue fringe.

Note that the AStar class must be subclassed for a specific problem domain!"""

import weakref
import sys
import traceback

import fibheap


##############################
#
# AStarNode
#
# A node in the A* search tree.  Stores heuristic value (h), cost-from-start
# value (g), overall cost value (f=h+g), a node id, an associated search state,
# a parent node, and a list of children
#
class AStarNode:
	"""A node in the A* search tree.  Stores heuristic value (h),
	cost-from-start value (g), overall cost value (f=h+g), a node id, an
	associated search state, a parent node, an optional piece of data
	associated with the edge, and a list of children.  Users'
	interaction with this type should be fairly limited, the only necessary
	direct uses being:
	
	-Storing a reference to an AStarNode associated with a given search state,
	for tracking visited states (see documentation on AStar for overriding
	its visit(), visited_state_node(), and clear_visited() methods).
	-Inspecting nodes along the final path from start to goal, as returned by
	AStar's result_path() method.
	
	To retrieve the state stored in an AStarNode, you can access its "state"
	data member."""
	
	def __init__(self, f=0, h=0, g=0, nid=0, state=None,
		     parent=None, parentedge=None,
		     children=None):
		self.f = f
		self.h = h
		self.g = g
		self.nid = nid
		self.state = state
		# Since parent holds a ref to children, we want a weak reference
		# to the parent to ensure garbage collection despite circular refs
		self.parent = weakref.proxy(parent) if parent else None
		self.parentedge = parentedge
		# Can't do children=[] in parameter list as all nodes wind up with
		# the same list
		self.children = children if children else []
		




#########################
# AStar
#
#
# An A* search problem instance.
#
# This is a generic class definition.  You should subclass it and override
# the is_goal(), successors(), clear_visited(), visit(), visited_state_node(),
# and heuristic() methods, as well as defining some kind of state class that
# represents a state in your search problem.
#
# Do call this class's __init__() in your subclass's, but make sure you can
# handle calls to clear_visited() and visit() before you do.
#
class AStar:
	"""An A* search problem instance.  A* search proceeds over states in the
	search space and stores visited states in tree nodes represented as
	AStarNode instances.  State types are defined by the user, but the user
	must be able to compare newly-generated states to old states to avoid
	revisiting states.  This class is intended to be subclassed by the user for
	a specific search space.  The following methods should be overridden:
	
	--is_goal()
	--successors()
	--clear_visited()
	--visit()
	--visited_state_node()
	--heuristic()
	
	NB:  The AStar.__init__() constructor should definitely be called by the
	subclass, but be prepared for it to call clear_visited() and visit()!  This
	means that whatever structure you need for tracking visited states
	should already be in place or should be initialized if necessary by those
	methods.
	
	After initialization, you can call search_step() to perform one step of the
	A* search, or search() to do a whole search to completion or failure.  Either
	will return True if a goal has been reached and False otherwise.
	
	After a successful search, you may inspect the data member "goal" to see
	the goal state and call result_path() for a list of AStarNode objects
	from start to goal.  You may also call num_nodes() for the total number
	of nodes in the search tree."""
	
	# Initialize with given start state.  If testGoalOnGeneration is set True,
	# newly-generated successor states are immediately tested for goal
	# conditions.  This may end search quicker but with a sub-optimal solution.
	def __init__(self, state, testGoalOnGeneration=False):
		"""Initializes a search with the given start state (a user-defined
		type).  If testGoalOnGeneration is True, then goal conditions will be
		tested on new states as soon as they are generated by the successors()
		method.  This may end the search more quickly but may also give
		sub-optimal results."""
		
		self.testGoalOnGeneration = testGoalOnGeneration
		self.set_start(state)
	
	# Set given state as the search's start state and initialize search
	# parameters.	
	def set_start(self, start):
		"""Sets the given state as the start state for the search and
		initializes the search tree and parameters."""
		
		# Reset the list of visited states
		self.clear_visited()
		
		self.fringe = fibheap.FibonacciHeap()
		# We don't have a goal node yet
		self.goal = None
		# Start with node id 0
		self.nid = 0
		
		if self.testGoalOnGeneration and self.is_goal(start):
			self.goal = start
			self.pathStates.append(start)
			return
		
		# Set up a search-tree node for the start state
		self.root = self.add_successor(None, start, None, 0)
		
	# Run search, to completion (returns True) or failure (returns False)
	def search(self):
		"""Run a search to completion or failure.  Return True if a solution
		is found, or False on failure."""
		
		while not self.fringe.is_empty():
			res = self.search_step()
			if res:
				return True
		return False
	
	# Run one step of a search (fringe removal, successor generation, goal
	# testing, append to fringe).  Return True if goal encountered, False else.
	def search_step(self):
		"""Perform one step of an A* search.  Return True if a goal state is
		found, False otherwise.  A step is defined as removing one node from
		the fringe, generating its successors, testing the node or its
		successors for goal conditions (depending on the value of
		testGoalOnGeneration when the AStar object was created), and appending
		the successors to the priority-queue search fringe."""
		
		if self.fringe.is_empty():
			return False
		
		# Get the fringe node with the minimal key and extract our search state
		n = self.fringe.extract_minimum()[1]
		del n.heapNode
		
		# Check goal condition
		if (not self.testGoalOnGeneration) and self.is_goal(n.state):
				self.goal = n
				return True
		
		# Generate successors
		res = self.successors(n.state)
		successors = res[0]
		costs = res[1] if len(res) > 1 else [1]*len(successors)
		edges = res[2] if len(res) > 2 else [None]*len(successors)
		
		# Do some tests on the successors
		for succ, cost, edge in zip(successors, costs, edges):
			# Goal test if appropriate
			if self.testGoalOnGeneration and self.is_goal(succ):
					self.goal = succ
					return True
			
			# Check to see if the state has been visited before in the search
			visited = self.visited_state_node(succ)
			# State was visited before, check its cost
			if visited:
				# This is a longer path from the start to this state than
				# before, so ignore this successor
				if n.g + cost >= visited.g:
					continue
				# This is a shorter path to this state than before, so
				# adjust its value in the fringe
				else:
					# First, remove the prior instance from the fringe
					if hasattr(visited,'heapNode'):
						try:
							self.fringe.delete(visited.heapNode)
						except (Exception) as e:
							print "Deleting from fringe failed.  "\
							      "Heuristic inadmissible or inconsistent?"
							raise
						        #print e
						        #traceback.print_exc()
					# Second, add the new version of the state
					self.add_successor(n,succ,edge,cost)
			# This state hasn't been visited before, so just add it
			else:
				self.add_successor(n,succ,edge,cost)
				
		return False
		
	# Add a successor to the given parent.  The state is contained in a new
	# node with the given cost from parent to child.  The new child is returned.
	def add_successor(self, parent, state, parentedge, cost):
		"""An internal operation.  The user should NOT call this method.
		Adds a new node to the search tree with the given parent."""
		
		# Set up a new node as a child of the parent
		g = parent.g + cost if parent else 0
		h = self.heuristic(state)
		f = h+g
		child = AStarNode(f,h,g,self.nid,state,parent,parentedge)
		self.nid += 1
		if parent:
			parent.children.append(child)
		# Put it in the fringe
		# Tuples are sorted lexicographically, so nodes are sorted by
		#  f-score first, then h-score, then g-score.
		child.heapNode = self.fringe.insert( (f,h,g), child )
		# Mark this state visited and map it to the new search node
		self.visit(state, child)
		
		return child
		
	# Number of nodes in the search tree.
	def num_nodes(self):
		"""Returns the number of nodes in the search tree."""
		return 1 + self.num_descendents(self.root)
		
	# Number of descendents from the given search-tree node.
	def num_descendents(self, node):
		"""Internal function.  Returns the number of descendents of the given
		search-tree node.  Users can call but there isn't much reason to."""
		count = len(node.children)
		for child in node.children:
			count += self.num_descendents(child)
		return count

	# Path from start to goal
	def result_path(self):
		"""Returns a list of AStarNode's up to the goal"""
		n = self.goal
		path = []
		while n != None:
			path.append(n)
			n = n.parent
		path.reverse()
		return path

	def result_cost(self):
		return self.goal.g

	def result_path_states(self):
		"""Returns a list of states up to the goal"""
		return [n.state for n in self.result_path()]

	def result_path_edges(self):
		"""Returns a list of edges up to the goal"""
		return [n.parentedge for n in self.result_path()]
		
		
		
	##################################
	# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	#
	# The following should be overloaded by a subclass!
	#
	
	# True if state is a goal state, otherwise False.  Users have no reason
	# to call this method.
	def is_goal(self, state):
		"""Returns True if the given state is a goal, otherwise False.  
		Internal; users may call this method but have little reason to.
		Abstract in this definition; should be overloaded by a subclass."""
		pass
		
	# Returns a pair of lists, one containing successor states and the other
	# containing costs to reach them.  Users have no reason to call this method.
	def successors(self, state):
		"""Returns a tuple containing one, two, or three lists:
		- A list of all legal successor states to the given state.
		- A list of corresponding costs to move to those states.
		- A list of data to be attached to each edge (e.g., an action
		  identifier)
		Internal; users may call this method but have little reason to.
		Abstract in this definition; should be overloaded by a subclass."""
		pass
		
	# Clears the current list of visited states.  Users should NOT call this
	# method.
	def clear_visited(self):
		"""Clears the list of visited states and their mappings to AStarNode
		instances.  See visit() and visited_state_node().  Internal; users
		should NEVER call this method.  Abstract in this
		definition; should be overloaded by a subclass."""
		pass
		
	# Marks a given state as visited and pointing to the given tree node.  Users
	# should NOT call this method.
	def visit(self, state, node):
		"""Marks the given state as visited and maps it to a given AStarNode.
		See visited_state_node() and clear_visited().  Internal; users should
		NEVER call this method.  Abstract in this
		definition; should be overloaded by a subclass."""
		pass
		
	# Returns the tree node corresponding to the given state if it has been
	# visited, otherwise None.  Users have no reason to call this method.
	def visited_state_node(self, state):
		"""If the given state has been visited, returns the associated
		AStarNode.  Otherwise, returns None.  See visit() and clear_visited().
		Internal; users may call this method but have little reason to.
		Abstract in this definition; should be overloaded by a subclass."""
		pass
		
	# Returns a heuristic value for the state.  Users have no reason to call
	# this method.
	def heuristic(self, state):
		"""Returns a heuristic value for the given state.  This heuristic
		should ideally fulfill the requirements for admissibility and
		consistency.  Internal; users may call this method but have little
		reason to.  Abstract in this definition; should be overloaded by
		a subclass."""
		return 0

	def tree_size(self,root=None):
		"""Returns the number of nodes in the search tree or subtree"""
		if root==None:
			if self.root == None: return 0
			return self.tree_size(self.root)
		return 1+sum(self.tree_size(c) for c in root.children)

	def best_node(self,quantity='f'):
		"""Upon a failed plan, returns the best node for the given
		quantity, which can be either f, g, or h."""
		q = [self.root]
		best = self.root
		bestval = self.root.f if quantity == 'f' else -self.root.g if quantity == 'g' else self.root.h
		while len(q) > 0:
			n = q.pop()
			cmpval = n.f if quantity == 'f' else -n.g if quantity == 'g' else n.h
			if cmpval < bestval:
				best = n
				bestval = cmpval
			for c in n.children:
				q.append(c)
		return best
